package handlers

import (
	"log"
	"npxy-go/pkg/tcp"
	"npxy-go/pkg/util/converter"
)

const ProtocolHeadSize = 6
const magicNumFixedLength uint16 = 10086

// LengthHeadHandler use to split or merge tcp frame
type LengthHeadHandler struct {
	tcp.BaseHandler
	merge *tcp.Buffer

	//total bytes to merge
	total int

	//the first two bytes are magic number that is 10086, rest 4 are length
	head []byte
}

func NewLengthHeadHandler() *LengthHeadHandler {
	return &LengthHeadHandler{
		head: make([]byte, 0, ProtocolHeadSize),
	}
}

func (l *LengthHeadHandler) OnMessage(ctx *tcp.SessionContext, buf *tcp.Buffer) {
	nextHandler := ctx.NextHandler()

	if nextHandler == nil {
		return
	}

	//in merging
	if l.inMerging() {
		l.handleMerging(ctx, buf)
		return
	}

	h, length, ok := l.getHead(buf)

	//bytes are not sufficient to hold head and length
	if !ok {
		return
	}

	//not an expected conn
	if h != magicNumFixedLength {
		log.Println("refused,not an expected conn:", ctx.String())
		ctx.Close()
		return
	}

	//normal
	if buf.FrameLen() == length {
		ctx.NextOnMessage(buf)
		return
	}

	//split
	if buf.FrameLen() > length {
		b1, b2 := buf.Split(length)
		inboundAndResetContext(ctx, b1)
		l.OnMessage(ctx, b2)
		return
	}

	//merge
	if buf.FrameLen() < length {
		l.prepareForMerging(length, buf)
	}
}

func (l *LengthHeadHandler) Write(ctx *tcp.SessionContext, buf *tcp.Buffer) {
	length := int32(buf.FrameLen())

	//add protocol head
	buf.AddInt32(length)
	buf.AddUInt16(magicNumFixedLength)

	ctx.PreWrite(buf)
}

func (l *LengthHeadHandler) inMerging() bool {
	return l.total > 0
}

func (l *LengthHeadHandler) getHead(buf *tcp.Buffer) (uint16, int, bool) {
	fl := buf.FrameLen()
	if fl == 0 {
		return 0, 0, false
	}

	rest := ProtocolHeadSize - len(l.head)

	if fl >= rest {
		l.head = append(l.head, buf.ReadBytes(rest)...)
		flag := converter.BytesToUInt16(l.head[:2])
		count := converter.BytesToInt32(l.head[2:])
		l.head = l.head[:0]

		return flag, int(count), true

	} else {
		l.head = append(l.head, buf.Frame()...)
		return 0, 0, false
	}
}

func (l *LengthHeadHandler) prepareForMerging(messageLen int, buf *tcp.Buffer) {
	l.total = messageLen
	l.allocMergeIfNecessary(messageLen)
	l.merge.AppendBytes(buf.Frame())
}

func (l *LengthHeadHandler) handleMerging(ctx *tcp.SessionContext, buf *tcp.Buffer) {
	rest := l.total - l.merge.FrameLen()
	fl := buf.FrameLen()
	if rest <= fl {
		l.merge.AppendBytes(buf.ReadBytes(rest))
		inboundAndResetContext(ctx, l.merge)
		l.finishMerging()

		// another message package, recursion

		if rest < fl {
			l.OnMessage(ctx, buf)
		}

	} else {
		l.merge.AppendBytes(buf.Frame())
	}
}

func (l *LengthHeadHandler) finishMerging() {
	l.total = 0
}

func (l *LengthHeadHandler) allocMergeIfNecessary(length int) {
	if l.merge == nil || l.merge.BodyCap() < length {
		l.merge = tcp.NewBufferWithHead(tcp.DefaultHeadSize, length)
	} else {
		l.merge.Reset()
	}
}

func inboundAndResetContext(ctx *tcp.SessionContext, buf *tcp.Buffer) {
	node := ctx.CurrentNode()
	ctx.NextOnMessage(buf)
	ctx.SetCurrentNode(node)
}
